package com.small.xx.ai.spring.ollama;

import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.retry.TransientAiException;
import org.springframework.retry.RetryCallback;
import org.springframework.retry.RetryContext;
import org.springframework.retry.RetryListener;
import org.springframework.retry.support.RetryTemplate;

import java.time.Duration;

public class OllamaChatBase {
    public String QWEN34B = "qwen3:4b";
    public String QWEN2Z5B = "qwen2:0.5b";
    // Ollama 服务地址
    public OllamaApi ollamaApi = OllamaApi.builder().baseUrl("http://127.0.0.1:11434").build();
    // ollama 模型  qwen3:4b
    public OllamaOptions ollamaOptions = OllamaOptions.builder().model(QWEN34B).build();
   // ollama RetryTemplate
    public RetryTemplate retryTemplate =
        RetryTemplate.builder().maxAttempts(1).retryOn(TransientAiException.class).fixedBackoff(Duration.ofSeconds(1))
            .withListener(new RetryListener() {
                @Override
                public <T extends Object, E extends Throwable> void onError(RetryContext context,
                    RetryCallback<T, E> callback, Throwable throwable) {
                    System.out.println("Retry error. Retry count:" + context.getRetryCount() + throwable);
                }
            }).build();
    // ollama chatModel
    public OllamaChatModel chatModel =
        OllamaChatModel.builder().ollamaApi(ollamaApi).defaultOptions(ollamaOptions).retryTemplate(retryTemplate)
            .build();

}
